from pypatnlp import *
import os
import sys
import multiprocessing

TASK_PATH = os.path.join('tasks', 'autorituvastus')
DATA_PATH = os.path.join(TASK_PATH, 'data')
REFERENCE_PATH = os.path.join(TASK_PATH, 'reference')
TRAINTEST_PATH = os.path.join(TASK_PATH, 'traintest')
AI_PATH = os.path.join(TASK_PATH, 'ai')
JOB_PATH = os.path.join(TASK_PATH, 'jobs')

NUM_THREADS = 20

radius=2
treshold=0.05
significance_treshold=0.025

JOB_PREFIX_TEMPLATE = '#PBS -N TRAIN_{0}'

JOB_TEMPLATE = '''

#PBS -l walltime=5:00:00
#PBS -l nodes=1:ppn=1
#PBS -l vmem=25g
#PBS -M timo_p@ut.ee
#PBS -d /home/petmanson/patnlp

# set the path
PATH=/home/petmanson/local/bin:/home/petmanson/local:/home/petmanson/pypfe/bin:/home/petmanson/local/jython:/home/petmanson/pfe/data/t3mesta:${PATH}
export PATH

# Some include paths for C++ compiler
CPLUS_INCLUDE_PATH=/home/petmanson/local/include:${CPLUS_INCLUDE_PATH}

export CPLUS_INCLUDE_PATH

# C compiler include path
C_INCLUDE_PATH=/home/petmanson/local/include:${C_INCLUDE_PATH}
export C_INCLUDE_PATH

# library path environment variable
LIBRARY_PATH=/home/petmanson/local/lib:${LIBRARY_PATH}
export LIBRARY_PATH

# Some library paths
LD_LIBRARY_PATH=/home/petmanson/local/lib:/usr/lib:/usr/local/lib:/home/petmanson/patnlp/lib:${LD_LIBRARY_PATH}
export LD_LIBRARY_PATH

# Python2.6 on vares NB! chance it to use it on krokodill   
#alias python='/www/apps/stacc-nrs/local/bin/python2.6'
#alias python26='/wrk/stacc-delfi/service/local/bin/python2.6'
#alias R='/wrk/stacc-delfi/software/bin/R'

# shared libraries included by Python
#export LD_LIBRARY_PATH=/www/apps/stacc-nrs/local/lib

# keeletehnoloogia vahendid
#export PATH=$PATH:/home/petmanson/keeletehnoloogia/bin:/home/petmanson/keeletehnoloogia/share/etmrf

PYTHONPATH=/home/petmanson:/home/petmanson/patnlp/:/home/petmanson/patnlp/lib:${PYTHONPATH}
export PYTHONPATH

JAVA_HOME=/usr/lib/jvm/java-openjdk/
export JAVA_HOME

JYTHONPATH=/home/timo/estner
export JYTHONPATH
'''

JOB_SUFFIX_TEMPLATE = 'python tasks/autorituvastus/test.py {0} {1} {2}'

def get_paths():
    return [name for name in os.listdir(DATA_PATH) if name.endswith('.pycorp')]

def get_traintest_paths():
    data = []
    for path in get_paths():
        trainpath = os.path.join(TRAINTEST_PATH, 'train_' + path)
        testpath  = os.path.join(TRAINTEST_PATH, 'test_' + path)
        name = os.path.basename(path)[:-7]
        data.append((name, trainpath, testpath))
    return data

def make_reference():
    ai = OneClassAI(radius=radius,
                    treshold=treshold,
                    significance_treshold = significance_treshold,
                    reference_dir = REFERENCE_PATH)
    ai.make_reference(REFERENCE_PATH)

def do_train_single(name, trainpath, testpath):
    name, trainpath, testpath = args
    traincorp = PyCorpus(trainpath, readonly=True)
    testcorp  = PyCorpus(testpath, readonly=True)
    
    ai = OneClassAI(radius=radius,
                    treshold=treshold,
                    significance_treshold = significance_treshold,
                    reference_dir = REFERENCE_PATH)
    try:
        ai.fit(traincorp)
        ys = ai.predict(testcorp)
        save_obj(ai, os.path.join(AI_PATH, name + '_ai.obj'))
    except RuntimeError, e:
        print e
    return 0

def make_jobs():
    f = open(os.path.join(JOB_PATH, 'submit.sh'), 'wb')
    for idx, (name, trainpath, testpath) in enumerate(get_traintest_paths()):
        prefix = JOB_PREFIX_TEMPLATE.format(idx)
        suffix = JOB_SUFFIX_TEMPLATE.format(name, trainpath, testpath)
        desc = prefix + JOB_TEMPLATE + suffix
        
        open(os.path.join(JOB_PATH, str(idx) + '.job'), 'wb').write(desc)
        f.write('qsub {0}.job\n'.format(idx))
    f.close()
z
if __name__ == '__main__':
    args = sys.argv[1:]
    if len(args) == 3:
        do_train_single(args[0], args[1], args[2])
    else:
        make_jobs()
